#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "rknn_model_runner.hpp"
#include "rknn_api.h"
#include "im2d.h"
#include "rk_mpi_mmz.h"

typedef struct DrmObject {
  int drm_buffer_fd;
  int drm_buffer_handle;
  size_t actual_size;
  uint8_t *drm_buf;
  uint64_t phy_addr;  // 物理地址
} DrmObject;

// void dump_tensor_attr(rknn_tensor_attr* attr)
// {
//   printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
//       "zp=%d, scale=%f\n",
//       attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
//       attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
//       get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
// }

void print_usage(const char* name)
{
    printf("Usage: %s <model_path> [options]\n", name);
    printf("Options:\n");
    printf("  -i <image_path>    Specify image file for inference\n");
    printf("  -c <camera_index>  Specify camera index for inference (default: 0)\n");
    printf("  -n <loop_count>    Specify inference loop count (default: 1)\n");
    printf("  -h                 Show this help message\n");
}

int main(int argc, char* argv[])
{
    if (argc < 2) {
        print_usage(argv[0]);
        return -1;
    }
    
    char* model_path = argv[1];
    char* image_path = NULL;
    int camera_index = -1;
    int loop_count = 1;
    int opt;
    
    // 解析命令行参数
    while ((opt = getopt(argc - 1, argv + 1, "i:c:n:h")) != -1) {
        switch (opt) {
            case 'i':
                image_path = optarg;
                break;
            case 'c':
                camera_index = atoi(optarg);
                break;
            case 'n':
                loop_count = atoi(optarg);
                break;
            case 'h':
                print_usage(argv[0]);
                return 0;
            default:
                print_usage(argv[0]);
                return -1;
        }
    }
    
    if (image_path == NULL && camera_index == -1) {
        printf("Either image path or camera index must be provided!\n");
        print_usage(argv[0]);
        return -1;
    }
    
    // 创建RKNN模型运行器
    RknnModelRunner* runner = new RknnModelRunner(model_path);
    if (!runner->is_model_loaded()) {
        printf("Failed to load model %s\n", model_path);
        delete runner;
        return -1;
    }
    
    int ret = 0;
    
    // 执行推理
    if (image_path) {
        printf("Running inference on image file: %s\n", image_path);
        // ret = runner->inference_from_file(image_path, loop_count);
    } else {
        printf("Running inference from camera: %d\n", camera_index);
        ret = runner->inference_from_camera(camera_index, loop_count);
    }
    
    // 清理资源
    delete runner;
    
    if (ret == 0) {
        printf("Inference completed successfully!\n");
    } else {
        printf("Inference failed with error code: %d\n", ret);
    }
    
    return ret;
}